package cn.showcon.firstapp.web;

import cn.showcon.firstapp.service.LlmService;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.ollama.OllamaChatModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.model.qianfan.QianfanChatModel;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.UserMessage;
import dev.langchain4j.service.V;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import dev.langchain4j.store.memory.chat.InMemoryChatMemoryStore;
import static dev.langchain4j.data.message.ChatMessageDeserializer.messagesFromJson;
import static dev.langchain4j.data.message.ChatMessageSerializer.messagesToJson;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;


import java.util.ArrayList;
import java.util.List;

/**
 * @Author Xue Lanbin
 */
@RestController
@RequestMapping("/langchain/test")
public class TestLangChainController {

    @Autowired
    private LlmService llmService;

    @GetMapping(value = "/simple-chat")
    public String simpleChat(String message) {
        try {
            QianfanChatModel qianfan = QianfanChatModel.builder()
                    .modelName("ERNIE-4.0-8K")
                    .endpoint("completions_pro")
                    .apiKey("jbug222LxcOwvHiWLuCOmjd2")
                    .secretKey("hktVnuX6XR7Cpeja0UClboehd2h8kEYp")
                    .responseFormat("json_object")
                    .build();
            String result = qianfan.generate(message);
            System.out.println("Qianfan: " + result);

        } catch (Exception e) {
            e.printStackTrace();
        }

        try {
            OllamaChatModel ollamaChatModel = OllamaChatModel.builder()
                    .baseUrl("http://192.168.1.240:11434")
                    .modelName("qwen2:7b")
                    .build();
            String r = ollamaChatModel.generate(message);
            System.out.println("Local:" +  r);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return "ok";
    }

    @GetMapping(value = "/chat")
    public String chat(String message) {
//        OllamaChatModel ollamaChatModel = OllamaChatModel.builder()
//                .baseUrl("http://192.168.1.240:11434")
//                .modelName("qwen2:7b")
//                .logRequests(true)
//                .logResponses(true)
//                .format("json")
//                .build();

        ChatLanguageModel chatLanguageModel = llmService.getModel("json");

        ChatMemoryStore chatMemoryStore = new InMemoryChatMemoryStore();
        ChatMemory chatMemory = MessageWindowChatMemory.builder()
                .maxMessages(10)
                .chatMemoryStore(chatMemoryStore)
                .build();

        Joke joke = AiServices.builder(Joke.class)
                .chatLanguageModel(chatLanguageModel)
                .chatMemory(chatMemory)
                .systemMessageProvider(chatId -> "")
                .build();

//        List<ChatMessage> messageList = new ArrayList<>();
//        messageList.add(SystemMessage.from("你负责讲关于旅行的笑话"));
//        messageList.add(UserMessage.from(message));
//        System.out.println(messagesToJson(messageList));

        System.out.println(joke.say("运动", 2));
        System.out.println(joke.say("旅行", 2));

//        Response<AiMessage> response = ollamaChatModel.generate(messageList);
//        return response.content().text();
        return "ok";
    }

    interface Joke {
        @dev.langchain4j.service.SystemMessage("你负责讲关于旅行的笑话")
        @UserMessage(value= """
                讲一个笑话，要求笑话是关于{{topic}}的。请一次性讲{{num}}条笑话。
                """)
        String say(@V("topic") String topic, @V("num") Integer num);
    }
}
